import os
import argparse
from collections import defaultdict

def get_dir_size(path, include_files=True):
    """计算目录大小（字节）"""
    total_size = 0
    for dirpath, _, filenames in os.walk(path):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            # 跳过符号链接
            if not os.path.islink(fp):
                total_size += os.path.getsize(fp)
    return total_size

def format_size(size_bytes):
    """将字节转换为易读格式"""
    units = ['B', 'KB', 'MB', 'GB', 'TB']
    size = size_bytes
    unit_index = 0
    
    while size >= 1024 and unit_index < len(units) - 1:
        size /= 1024
        unit_index += 1
    
    return f"{size:.2f} {units[unit_index]}"

def analyze_directory(root_dir, depth=1):
    """分析目录结构和大小"""
    try:
        print(f"分析目录: {root_dir}")
        print(f"最大深度: {depth}")
        
        # 计算总大小
        total_size = get_dir_size(root_dir)
        print(f"总大小: {format_size(total_size)}\n")
        
        # 收集子目录大小
        dir_sizes = defaultdict(int)
        
        # 遍历一级子目录
        for entry in os.listdir(root_dir):
            entry_path = os.path.join(root_dir, entry)
            if os.path.isdir(entry_path):
                dir_sizes[entry] = get_dir_size(entry_path)
        
        # 按大小排序（降序）
        sorted_dirs = sorted(dir_sizes.items(), key=lambda x: x[1], reverse=True)
        
        # 显示结果
        print("子目录大小统计:")
        for dir_name, size in sorted_dirs:
            percentage = (size / total_size) * 100 if total_size > 0 else 0
            print(f"  {dir_name}: {format_size(size)} ({percentage:.1f}%)")
        
        # 如果需要更深层次的分析，可以递归调用
        # if depth > 1:
        #     for dir_name, _ in sorted_dirs[:5]:  # 分析前5个最大的目录
        #         print(f"\n分析 {dir_name}:")
        #         analyze_directory(os.path.join(root_dir, dir_name), depth-1)
                
    except Exception as e:
        print(f"分析失败: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='文件夹大小分析脚本')
    parser.add_argument('directory', help='要分析的目录路径')
    parser.add_argument('--depth', type=int, default=1, help='分析深度，默认1级')
    
    args = parser.parse_args()
    analyze_directory(args.directory, args.depth)
